#include "jassvm_extension.h"
#include "Extern_global_Variable.h"

String GetCodeName2(PInstance Instance, PInstruction Code)
{
	String Result;
	PInstruction FunctionLabel;

	Result = new char[100];

	if (Code != NULL)
	{
		FunctionLabel = Code - 1;
		if (TOperationCode((*FunctionLabel).id) == opFunction)
		{
			sprintf(Result,"%s", (*Instance).Script->Strings->GetItem((*FunctionLabel).Parameter));
		}
		else
		{
			sprintf(Result,"0x%x",Code);
		};
	}
	else
	{
		sprintf(Result,"nullptr");
	};
	return Result;
};

TScriptHandle FindScriptHandle(PJassScript Script)
{
	TScriptHandle Result;
	Dword vvi;
	PJassThreadLocal Environment;

	Result = 0;

	Environment = GetJassThreadLocal();
	for (vvi = 0; vvi <= (*Environment).Scripts.GetLength() - 1; vvi++)
	{
		if (Script == (*Environment).Scripts.GetItem(vvi))
		{
			Result = vvi;
			break;
		};
	};
	return Result;
};

TInstanceHandle hook_RegisterInstance2() // fastcall
{
	TInstanceHandle Result; 
	PInstance Instance;

	__asm
	{
		mov Instance, ecx
	};

	Result = orig_RegisterInstance();
	InstanceToScriptTable.InsertValue(Result, FindScriptHandle((*Instance).Script));
	return Result;
};

procedure InitScriptFuncTable(TScriptHandle hScript)
{
	Integer vvi;
	PJassThreadLocal Environment;
	TFuncTable Table;
	PJassScript Script;
	PInstruction Instruction;
	TCodeIterator* Instructions;
	Dword CurrentFunction = 0;
	PNativeFunc Native;
	Dword StringIndex;

	Environment = GetJassThreadLocal();
	Script = (*Environment).Scripts.GetItem(hScript);
	Instructions = new TCodeIterator(Script);


	Table = new TFuncData[(*Script).Strings->GetLength()];

	for (vvi = 0; vvi <= Instructions->FInstructionsCount - 1; vvi++)
	{
		Instruction = Instructions->GetCurrent();
		Instructions->MoveNext();
		switch (TOperationCode((*Instruction).id))
		{
			case opEndProgram: goto m1; break;
			case opFunction: CurrentFunction = (*Instruction).Parameter; break;
			case opArgument: Table[CurrentFunction].ArgsCount++; break;
		};
	};
m1:
	for (vvi = 0; vvi <= (*Environment).Natives.Items.GetLength() - 1; vvi++)
	{
		if (vvi == 0)
		{
			Native = (*Environment).Natives.Items.GetFirst();
		}
		else
		{
			Native = (*Environment).Natives.Items.GetNextItem(Native);
		}
		if ((*Script).Strings->GetItemIndex(
			(*Native).HashKeyData.Key,
			StringIndex
		))
		{
			Table[StringIndex].ArgsCount = (*Native).ArgsCount;
		};
	};

	ScriptToFuncTable.InsertValue(hScript,Table);
};

TScriptHandle __stdcall hook_Compile(Dword ErrorCallbackParam, Dword RiskLevel) // fastcall
{
	TScriptHandle Result;
	Result = orig_Compile(ErrorCallbackParam, RiskLevel);
	InitScriptFuncTable(Result);
	return Result;
};

procedure jassvm_extension_OnInit()
{
	char vStr[100];
	if (!Hook(
		Pointer(&_RegisterInstance),
		&hook_RegisterInstance2,
		Pointer(&orig_RegisterInstance),
		NULL
	))
	{
		sprintf(vStr,"Failed to hook RegisterInstance.");
		MessageBoxA(NULL,vStr,"Error!",MB_OK | MB_ICONWARNING);
	}

	if (!Hook(
		Pointer(_Compile),
		&hook_Compile,
		Pointer(orig_Compile),
		NULL
	))
	{
		sprintf(vStr,"Failed to hook compiler.");
		MessageBoxA(NULL,vStr,"Error!",MB_OK | MB_ICONWARNING);
	}
};

TFuncTable GetScriptFunctions(TScriptHandle hScript)
{
	return ScriptToFuncTable.GetValue(hScript);
};

TScriptHandle GetInstanceScriptHandle(PInstance Instance)
{
	return InstanceToScriptTable.GetValue((*Instance).Handle);
};

PInstruction TCodeIterator::GetCurrent()
{
	return &(*FScript).Code[FCurrentIndex];
};

Boolean TCodeIterator::MoveNext()
{
	Boolean Result;
	if (FCurrentIndex < (FInstructionsCount - 1))
	{
		FCurrentIndex++;
		Result = TRUE;
	}
	else
	{
		Result = FALSE;
	};
	return Result;
};

procedure TCodeIterator::Reset()
{
	FCurrentIndex = -1;
};

TCodeIterator::TCodeIterator(PJassScript AScript)
{
	FScript = AScript;
	FInstructionsCount = (*FScript).CodeSize / sizeof(TInstruction);
	Reset();
};

TCodeIterator TCodeIterator::GetEnumerator()
{
	return *this;
};

void jassvm_extension_main()
{
	InstanceToScriptTable.THandleTableM::THandleTableM();
	ScriptToFuncTable.THandleTableM::THandleTableM();
	wc3main_SubscribeOnInit(&jassvm_extension_OnInit);
}

